# 金肯特纳交易策略
# 基于典型价格(最高价+最低价+收盘价)/3的移动平均线通道，配合真实波动幅度(ATR)构建交易信号
def init(context):
    # 初始化账户
    set_subportfolios([{"cash": 0, "type": 'stock'}, {"cash": 100000, "type": "future"}])
    log.info('金肯特纳策略开始运行')
    
    # 设置交易参数
    set_commission(PerShare(type='future', cost=0.000045))
    set_slippage(PriceSlippage(0.005), 'future')
    set_margin_rate('RB', 0.08, 0.09)
    set_volume_limit(0.25, 0.5)
    
    # 设置交易标的
    context.ins = 'RB9999'
    set_benchmark(context.ins)
    subscribe(context.ins)
    
    # 策略参数
    context.ma_period = 56      # 均线周期
    context.atr_period = 3      # ATR周期
    context.trade_size = 20     # 交易手数
    
    # 初始化数据容器
    context.data = {
        'highs': [],
        'lows': [],
        'closes': [],
    }

def handle_bar(context, bar_dict):
    # 获取当前K线数据
    symbol = context.ins
    high = bar_dict[symbol].high
    low = bar_dict[symbol].low
    close = bar_dict[symbol].close
    
    # 保存历史数据
    context.data['highs'].append(high)
    context.data['lows'].append(low)
    context.data['closes'].append(close)
    
    # 计算典型价格(最高+最低+收盘)/3
    typ_price = (high + low + close) / 3.0
    
    # 计算真实波幅(TR)
    if len(context.data['closes']) < 2:
        tr = high - low  # 首日使用普通波幅
    else:
        prev_close = context.data['closes'][-2]  # 前一日收盘价
        tr = max(high - low, 
                abs(high - prev_close), 
                abs(low - prev_close))
    
    # 检查数据是否足够计算指标
    if len(context.data['closes']) < max(context.ma_period, context.atr_period):
        return
    
    # 获取历史数据计算指标
    # 计算典型价格序列
    typ_prices = [
        (context.data['highs'][i] + context.data['lows'][i] + context.data['closes'][i]) / 3.0 
        for i in range(len(context.data['closes']))
    ]
    
    # 计算移动平均
    ma = sum(typ_prices[-context.ma_period:]) / context.ma_period
    
    # 计算ATR
    # 先计算所有TR值
    trs = []
    for i in range(len(context.data['closes'])):
        if i == 0:
            tr_val = context.data['highs'][i] - context.data['lows'][i]
        else:
            prev_close = context.data['closes'][i-1]
            high = context.data['highs'][i]
            low = context.data['lows'][i]
            tr_val = max(high - low, abs(high - prev_close), abs(low - prev_close))
        trs.append(tr_val)
    
    atr = sum(trs[-context.atr_period:]) / context.atr_period
    
    # 计算通道
    upper_band = ma + atr
    lower_band = ma - atr
    
    # 获取当前持仓
    future_account = context.portfolio.future_account
    positions = future_account.positions
    
    # 初始化持仓数量
    long_amount = 0
    short_amount = 0
    
    # 安全获取持仓信息
    if symbol in positions:
        position = positions[symbol]
        long_amount = position.long_amount
        short_amount = position.short_amount
    
    # 获取前一日收盘价位置（用于判断穿越）
    prev_close = context.data['closes'][-2] if len(context.data['closes']) >= 2 else None
    
    # 交易信号
    # 1. 上穿通道上轨 - 开多
    if prev_close is not None and close > upper_band and prev_close <= upper_band:
        if short_amount > 0:  # 先平空
            order_future(symbol, short_amount, 'close', 'short', None)
        if long_amount == 0:  # 再开多
            order_future(symbol, context.trade_size, 'open', 'long', None)
            log.info('价格突破上轨，开多仓')
    
    # 2. 下穿通道下轨 - 开空
    if prev_close is not None and close < lower_band and prev_close >= lower_band:
        if long_amount > 0:  # 先平多
            order_future(symbol, long_amount, 'close', 'long', None)
        if short_amount == 0:  # 再开空
            order_future(symbol, context.trade_size, 'open', 'short', None)
            log.info('价格突破下轨，开空仓')
    
    # 3. 价格跌破下轨 - 平多
    if close < lower_band and long_amount > 0:
        order_future(symbol, long_amount, 'close', 'long', None)
        log.info('价格跌破下轨，平多仓')
    
    # 4. 价格突破上轨 - 平空
    if close > upper_band and short_amount > 0:
        order_future(symbol, short_amount, 'close', 'short', None)
        log.info('价格突破上轨，平空仓')
    
    # 维护数据长度
    max_length = max(context.ma_period, context.atr_period) * 3
    for key in context.data:
        if len(context.data[key]) > max_length:
            context.data[key] = context.data[key][-max_length:]

# 其他必要函数
def before_trading(context):
    date = get_datetime().strftime('%Y-%m-%d %H:%M:%S')
    log.info('{} 盘前运行'.format(date))

def after_trading(context):
    time = get_datetime().strftime('%Y-%m-%d %H:%M:%S')
    log.info('{} 盘后运行'.format(time))
    log.info('一天结束')